import torch
from torch import nn

# 检查 CUDA 是否可用
print("CUDA is available:", torch.cuda.is_available())

# 打印可用的 GPU 数量
print("Number of GPUs available:", torch.cuda.device_count())

# 打印当前默认设备
print("Current device:", torch.cuda.current_device())

# 打印设备名称
if torch.cuda.is_available():
    print("Device name:", torch.cuda.get_device_name(0))